"""
自动权重多任务损失函数

主要功能：
1. 自动学习和平衡多个损失项的权重
2. 避免手动调整多任务学习中的权重超参数
3. 通过反向传播自动优化各损失的重要性

特点：
- 动态权重调整
- 端到端训练
- 无需人工干预
- 适用于任意数量的损失项
"""

import torch
import torch.nn as nn

class AutomaticWeightedLoss(nn.Module):
    """
    自动加权多任务损失
    
    原理：
    - 为每个损失项分配一个可学习的权重参数
    - 权重通过梯度下降自动优化
    - 使用权重的平方倒数作为实际权重，保证稳定性
    """
    def __init__(self, num=2):
        """
        初始化自动权重损失函数
        Args:
            num: int，损失项的数量（默认为2）
        说明：
            - 初始化每个损失项的权重为1
            - 权重参数可通过反向传播自动更新
        """
        super(AutomaticWeightedLoss, self).__init__()
        params = torch.ones(num, requires_grad=True)
        self.params = torch.nn.Parameter(params)

    def forward(self, *x):
        """
        计算加权后的总损失
        Args:
            *x: 损失项列表，如[loss1, loss2, ...]
        Returns:
            losses: 加权后的总损失
            
        实现原理：
        1. 对每个损失项应用其对应的权重
        2. 权重计算公式：loss_weighted = 1/(2*w^2)*loss + log(1+w^2)
        3. 第二项log(1+w^2)用于防止权重趋近于0
        """
        losses = 0
        for i, loss in enumerate(x[0]):
            # 使用权重的平方倒数作为实际权重
            weight = 0.5 / (self.params[i] ** 2)
            # 添加正则项防止权重过小
            reg_term = torch.log(1 + self.params[i] ** 2)
            # 计算加权损失
            weighted_loss = weight * loss + reg_term
            losses = losses + weighted_loss
        return losses

if __name__ == '__main__':
    awl = AutomaticWeightedLoss(2)
    print(awl.parameters())  # 查看初始权重参数